from langgraph.config import get_stream_writer
from langchain_core.messages import SystemMessage
from langchain_core.runnables import RunnableConfig
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import PromptTemplate
from datetime import datetime
from pydantic import BaseModel, Field

from app.graph.BaseNode import BaseNode
from app.graph.state import State
from app.graph.llm import chat_zhipu


class Intent(BaseModel):
    intent: str = Field(description="用户意图")


_parser = JsonOutputParser(pydantic_object=Intent)


# TODO: 先放一边

_pt = """# Role: 文本提取专家

## Goals

从以下对话中提取用户意图。如果匹配到下面列出意图，返回意图的 name，否则返回 unknown。下面是当前支持的意图：

- name: weather
  desc: 查看天气
- name: news
  desc: 查看新闻

应当输出 JSON 格式的结果。

## Context

以下是用户最近说过的话，最上面说的时间最早，最下面说得时间最晚。

{user_messages}

## Examples

{
    "intent": "weather"
}
或者
{
    "intent": "unknown"
}
"""

_prompt = PromptTemplate(template=_pt, input_variables=["user_messages"])


class NodeIntentRecog(BaseNode):
    def __call__(self, state: State, config: RunnableConfig):
        tid = self.get_tid(config)
        messages = self.get_latest_human_messages(tid)
        if len(messages) == 0:
            raise RuntimeError("没有用户回复")  # TODO: 测试抛出异常会怎样

        prompt = _prompt.invoke(
            {"user_messages": "\n".join([f"- {it.content}" for it in messages])}
        ).to_string()

        messages.append(SystemMessage(content=prompt))

        resp = chat_zhipu.invoke(messages)

        result = _parser.invoke(resp.content)

        return {"intent": result["intent"]}
